"""水文流量数据仓储层

提供对水文流量数据的数据访问操作，封装数据库查询逻辑。
"""

from datetime import datetime
from typing import List, Optional, Tuple, Dict, Any

from sqlalchemy import func, and_
from sqlalchemy.orm import Session

from ..models.hydrometric import HydrometricDataTimeseries, HydrometricDataRegular


class HydrometricRepository:
    """水文流量数据仓储类"""

    def __init__(self, session: Session):
        """初始化仓储
        
        Args:
            session: 数据库会话
        """
        self.session = session

    def get_station_count(self, table_type: str = 'timeseries') -> int:
        """获取测站数量
        
        Args:
            table_type: 表类型 ('timeseries' 或 'regular')
            
        Returns:
            测站数量
        """
        model = HydrometricDataTimeseries if table_type == 'timeseries' else HydrometricDataRegular
        return self.session.query(func.count(func.distinct(model.station_id))).scalar()

    def get_total_records(self, table_type: str = 'timeseries') -> int:
        """获取总记录数
        
        Args:
            table_type: 表类型 ('timeseries' 或 'regular')
            
        Returns:
            总记录数
        """
        model = HydrometricDataTimeseries if table_type == 'timeseries' else HydrometricDataRegular
        return self.session.query(func.count(model.id)).scalar()

    def get_date_range(self, table_type: str = 'timeseries') -> Tuple[Optional[datetime], Optional[datetime]]:
        """获取数据时间范围
        
        Args:
            table_type: 表类型 ('timeseries' 或 'regular')
            
        Returns:
            (最早时间, 最晚时间)
        """
        model = HydrometricDataTimeseries if table_type == 'timeseries' else HydrometricDataRegular
        result = self.session.query(
            func.min(model.timestamp),
            func.max(model.timestamp)
        ).first()
        return result if result else (None, None)

    def query_by_station_and_time(
            self,
            station_id: str,
            start_time: datetime,
            end_time: datetime,
            table_type: str = 'timeseries',
            limit: int = 10000
    ) -> List[Dict[str, Any]]:
        """按测站和时间范围查询数据
        
        Args:
            station_id: 测站ID
            start_time: 开始时间
            end_time: 结束时间
            table_type: 表类型 ('timeseries' 或 'regular')
            limit: 限制返回记录数，默认10000条
            
        Returns:
            查询结果列表（按时间倒序排列）
        """
        model = HydrometricDataTimeseries if table_type == 'timeseries' else HydrometricDataRegular

        results = self.session.query(model).filter(
            and_(
                model.station_id == station_id,
                model.timestamp >= start_time,
                model.timestamp <= end_time
            )
        ).order_by(model.timestamp.desc()).limit(limit).all()

        return [
            {
                'station_id': r.station_id,
                'timestamp': r.timestamp,
                'flow_rate': r.flow_rate
            }
            for r in results
        ]

    def aggregate_by_stations_and_time(
            self,
            station_ids: List[str],
            start_time: datetime,
            end_time: datetime,
            table_type: str = 'timeseries'
    ) -> List[Dict[str, Any]]:
        """按多个测站和时间范围聚合查询
        
        Args:
            station_ids: 测站ID列表
            start_time: 开始时间
            end_time: 结束时间
            table_type: 表类型 ('timeseries' 或 'regular')
            
        Returns:
            聚合结果列表
        """
        model = HydrometricDataTimeseries if table_type == 'timeseries' else HydrometricDataRegular

        results = self.session.query(
            model.station_id,
            func.avg(model.flow_rate).label('avg_flow'),
            func.max(model.flow_rate).label('max_flow'),
            func.min(model.flow_rate).label('min_flow'),
            func.count(model.flow_rate).label('count')
        ).filter(
            and_(
                model.station_id.in_(station_ids),
                model.timestamp >= start_time,
                model.timestamp <= end_time
            )
        ).group_by(model.station_id).order_by(model.station_id).all()

        return [
            {
                'station_id': r.station_id,
                'avg_flow': float(r.avg_flow) if r.avg_flow else 0.0,
                'max_flow': float(r.max_flow) if r.max_flow else 0.0,
                'min_flow': float(r.min_flow) if r.min_flow else 0.0,
                'count': r.count
            }
            for r in results
        ]

    def query_time_range_stats(
            self,
            start_time: datetime,
            end_time: datetime,
            table_type: str = 'timeseries'
    ) -> Dict[str, Any]:
        """查询时间范围内的统计信息
        
        Args:
            start_time: 开始时间
            end_time: 结束时间
            table_type: 表类型 ('timeseries' 或 'regular')
            
        Returns:
            统计信息字典
        """
        model = HydrometricDataTimeseries if table_type == 'timeseries' else HydrometricDataRegular

        result = self.session.query(
            func.count(model.id).label('total_records'),
            func.count(func.distinct(model.station_id)).label('station_count'),
            func.avg(model.flow_rate).label('avg_flow'),
            func.max(model.flow_rate).label('max_flow'),
            func.min(model.flow_rate).label('min_flow')
        ).filter(
            and_(
                model.timestamp >= start_time,
                model.timestamp <= end_time
            )
        ).first()

        return {
            'total_records': result.total_records if result else 0,
            'station_count': result.station_count if result else 0,
            'avg_flow': float(result.avg_flow) if result and result.avg_flow else 0.0,
            'max_flow': float(result.max_flow) if result and result.max_flow else 0.0,
            'min_flow': float(result.min_flow) if result and result.min_flow else 0.0
        }

    def clear_all_data(self, table_type: str = 'timeseries') -> int:
        """清空所有数据
        
        Args:
            table_type: 表类型 ('timeseries' 或 'regular')
            
        Returns:
            删除的记录数
        """
        model = HydrometricDataTimeseries if table_type == 'timeseries' else HydrometricDataRegular
        count = self.session.query(model).count()
        self.session.query(model).delete()
        self.session.commit()
        return count

    def bulk_insert(
            self,
            data: List[Dict[str, Any]],
            table_type: str = 'timeseries'
    ) -> int:
        """批量插入数据
        
        Args:
            data: 数据列表
            table_type: 表类型 ('timeseries' 或 'regular')
            
        Returns:
            插入的记录数
        """
        if not data:
            return 0

        model = HydrometricDataTimeseries if table_type == 'timeseries' else HydrometricDataRegular

        # 使用bulk_insert_mappings进行批量插入
        self.session.bulk_insert_mappings(model, data)
        self.session.commit()

        return len(data)

    def get_distinct_station_ids(self, limit: int = None, table_type: str = 'timeseries') -> List[str]:
        """获取不重复的测站ID列表
        
        Args:
            limit: 限制返回数量
            table_type: 表类型 ('timeseries' 或 'regular')
            
        Returns:
            测站ID列表
        """
        model = HydrometricDataTimeseries if table_type == 'timeseries' else HydrometricDataRegular
        query = self.session.query(func.distinct(model.station_id))
        
        if limit:
            query = query.limit(limit)
            
        results = query.all()
        return [result[0] for result in results]

    def query_by_stations_and_time(
            self,
            station_ids: List[str],
            start_time: datetime,
            end_time: datetime,
            table_type: str = 'timeseries',
            limit: int = 50000
    ) -> List[Dict[str, Any]]:
        """按多个测站和时间范围查询数据
        
        Args:
            station_ids: 测站ID列表
            start_time: 开始时间
            end_time: 结束时间
            table_type: 表类型 ('timeseries' 或 'regular')
            limit: 限制返回记录数，默认50000条
            
        Returns:
            查询结果列表（按时间倒序排列）
        """
        model = HydrometricDataTimeseries if table_type == 'timeseries' else HydrometricDataRegular

        results = self.session.query(model).filter(
            and_(
                model.station_id.in_(station_ids),
                model.timestamp >= start_time,
                model.timestamp <= end_time
            )
        ).order_by(model.timestamp.desc()).limit(limit).all()

        return [
            {
                'station_id': r.station_id,
                'timestamp': r.timestamp,
                'flow_rate': r.flow_rate
            }
            for r in results
        ]
